{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AbwmoJusmueU"
   },
   "source": [
    "# Reinforcement Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7xcuGNp-mueX"
   },
   "source": [
    "В этом задании постараемся разобраться в проблеме обучения с подкреплением, реализуем алгоритм REINFORCE и научим агента с помощью этого алгоритма играть в игру Cartpole."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4HoHMD_GmueZ"
   },
   "source": [
    "Установим и импортируем необходимые библиотеки, а также вспомогательные функции для визуализации игры агента."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kHKfskxa7kaI"
   },
   "outputs": [],
   "source": [
    "!pip install gym pandas torch matplotlib pyvirtualdisplay > /dev/null 2>&1\n",
    "!apt-get install -y xvfb python-opengl ffmpeg x11-utils > /dev/null 2>&1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WiX0JVM3z0Ru"
   },
   "outputs": [],
   "source": [
    "from IPython.display import clear_output, HTML\n",
    "from IPython import display as ipythondisplay\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IsjSHknYl1rj"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import glob\n",
    "import time\n",
    "import io\n",
    "import base64\n",
    "import gym\n",
    "from gym.wrappers import Monitor\n",
    "import torch\n",
    "import collections\n",
    "import pandas as pd\n",
    "from torch import nn\n",
    "from torch.optim import Adam\n",
    "from torch.distributions import Categorical"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "b5KfVGq1n_8s"
   },
   "outputs": [],
   "source": [
    "from pyvirtualdisplay import Display\n",
    "display = Display(visible=0, size=(1400, 900))\n",
    "display.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4xGh4QxDnlFw"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Utility functions to enable video recording of gym environment and displaying it\n",
    "To enable video, just do \"env = wrap_env(env)\"\"\n",
    "\"\"\"\n",
    "\n",
    "def show_video():\n",
    "    mp4list = glob.glob('video/*.mp4')\n",
    "    if len(mp4list) > 0:\n",
    "        mp4 = mp4list[0]\n",
    "        video = io.open(mp4, 'r+b').read()\n",
    "        encoded = base64.b64encode(video)\n",
    "        ipythondisplay.display(HTML(data='''<video alt=\"test\" autoplay \n",
    "                loop controls style=\"height: 400px;\">\n",
    "                <source src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\" />\n",
    "             </video>'''.format(encoded.decode('ascii'))))\n",
    "    else: \n",
    "        print(\"Could not find video\")\n",
    "    \n",
    "\n",
    "def wrap_env(env):\n",
    "    env = Monitor(env, './video', force=True)\n",
    "    return env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Dfi89tXjmuet"
   },
   "outputs": [],
   "source": [
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') #позволяет перенести тензор на GPU, если он доступен в системе"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AhDVTPxymuez"
   },
   "source": [
    "## OpenAI Gym"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VZb-IIq4muez"
   },
   "source": [
    "[OpenAI Gym](https://gym.openai.com) это набор сред для разработки и сравнения алгоритмов обучения с подкреплением."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "fFkuRjxNmue2"
   },
   "source": [
    "OpenAI Gym предоставляет простой и универсальный API ко многим средам с разными свойствами, как простым так и сложным:\n",
    "* Классические задачи управления и игрушечные примеры, которые можно найти в учебниках и на которых демонстрируется работа алгоритмов обучения с подкреплением (одна из этих сред используется в этом задании)\n",
    "* Игры Atari (оказали огромное влияние на достижения в обучении с подкреплением в последние годы)\n",
    "* 2D и 3D среды для контроля роботов в симуляции (используют проприетарный движок [Mojuco](http://www.mujoco.org))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NIu0Uhj7mue5"
   },
   "source": [
    "Рассмотрим, как устроена среда [CartPole-v0](https://gym.openai.com/envs/CartPole-v0), с которой мы будем работать."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KYHUeiMXmue6"
   },
   "source": [
    "Для этого создадим среду и выведем ее описание."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "b9BRqxoDmue7"
   },
   "outputs": [],
   "source": [
    "env = gym.make(\"CartPole-v0\")\n",
    "print(env.env.__doc__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tovAoT5dmufD"
   },
   "source": [
    "Из этого описания мы можем узнать, как устроены пространства состояний и действий в этой среды, какие награды получаются на каждом шаге, а также, что нам необходимо сделать, чтобы научиться \"решать\" эту среду, а именно достич средней награды больше 195.0 или больше за 100 последовательных запусков агента в этой среде. Именно такого агента мы и попробуем создать и обучить."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hZ98fYaxmufF"
   },
   "source": [
    "Но для начала напишем вспомогательную функцию, которая будет принимать на вход среду, агента и число эпизодов, и возвращать среднюю награду за 100 эпизодов. С помощью этой функции мы сможем протестировать, насколько хорошо обучился наш агент, а также визуализировать его поведение в среде."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0jXEOwd7njUm"
   },
   "outputs": [],
   "source": [
    "def test_agent(env, agent=None, n_episodes=100):\n",
    "    \"\"\"Runs agent for n_episodes in environment and calclates mean reward.\n",
    "    \n",
    "    Args:\n",
    "        env: The environment for agent to play in\n",
    "        agent: The agent to play with. Defaults to None - \n",
    "            in this case random agent is used.\n",
    "        n_episodes: Number of episodes to play. Defaults to 100.\n",
    "\n",
    "    Returns:\n",
    "        Mean reward for 100 episodes.\n",
    "    \"\"\"\n",
    "    total_reward = []\n",
    "    for episode in range(n_episodes):\n",
    "        episode_reward = 0\n",
    "        observation = env.reset()\n",
    "        t = 0\n",
    "        while True:\n",
    "            if agent:\n",
    "                with torch.no_grad():\n",
    "                    probs = agent(torch.FloatTensor(observation).to(device))\n",
    "                    dist = Categorical(probs)\n",
    "                    action = dist.sample().item()\n",
    "            else:\n",
    "                action = env.action_space.sample()\n",
    "            observation, reward, done, info = env.step(action)\n",
    "            episode_reward += reward\n",
    "            t += 1\n",
    "            if done:\n",
    "                print(\"Episode {} finished after {} timesteps\".format(episode+1, t+1))\n",
    "                break\n",
    "        total_reward.append(episode_reward)\n",
    "        env.close()\n",
    "                   \n",
    "    return np.mean(total_reward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "TqnYqFVSmufL"
   },
   "source": [
    "Протестируем и визуализируем случайного агента (параметр ```agent=False```)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yqVokBDFmufM",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "test_agent(env, agent=False, n_episodes=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8bzNtyllmufS"
   },
   "source": [
    "Как видно, наш случайный агент выступает не очень хорошо и в среднем может удержать шест всего около 20 шагов."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BNcZLQpzmufU"
   },
   "source": [
    "Напишем функцию для визуализации агента и посмотрим на случайного агента."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GgThoRPgoMcs"
   },
   "outputs": [],
   "source": [
    "def agent_viz(env=\"CartPole-v0\", agent=None):\n",
    "    \"\"\"Visualizes agent play in the given environment.\n",
    "    \n",
    "    Args:\n",
    "        env: The environment for agent to play in. Defaults to CartPole-v0.\n",
    "        agent: The agent to play with. Defaults to None - \n",
    "            in this case random agent is used.\n",
    "\n",
    "    Returns:\n",
    "        Nothing is returned. Visualization is created and can be showed\n",
    "            with show_video() function.\n",
    "    \"\"\"\n",
    "    env = wrap_env(gym.make(env))\n",
    "    observation = env.reset()\n",
    "    while True:\n",
    "        env.render() \n",
    "        if agent:\n",
    "            with torch.no_grad():\n",
    "                probs = agent(torch.FloatTensor(observation).to(device))\n",
    "                dist = Categorical(probs)\n",
    "                action = dist.sample().item()\n",
    "        else:\n",
    "            action = env.action_space.sample()\n",
    "        observation, reward, done, info = env.step(action)\n",
    "        if done:\n",
    "            break\n",
    "            \n",
    "    env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "b5FLllE_pHSf"
   },
   "outputs": [],
   "source": [
    "agent_viz()\n",
    "show_video()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hpQF0xgJmufW"
   },
   "source": [
    "Попробуем применить обучение с подкреплением и алгоритм REINFORCE для того, чтобы в среднем за 100 эпиздов мы держали шест не менее 195 шагов."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "M2SkO__WmufX"
   },
   "source": [
    "## REINFORCE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QW95huRpmufY"
   },
   "source": [
    "Вспомним, что из себя представляет алгоритм REINFORCE (Sutton & Barto) <img src=\"//i.imgur.com/bnASTrY.png\" width=\"700\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "CFD7oH9mmufZ"
   },
   "source": [
    "1. Инициализуем политику (в качестве политики мы будем использовать глубокую нейронную сеть).\n",
    "2. \"Играем\" в среде эпизод, используя нашу политику, или несколько (мы будем использовать последний вариант) и собираем данные о состояниях, действиях и полученных наградах. \n",
    "3. Для каждого состояния в собранных эпизодах вычисляем сумму дисконтированных наград, полученных из этого состояния, а также логорифм правдоподобия предпринятого действия в этом состоянии для нашей политики.\n",
    "4. Обновляем параметры нашей политики по формуле на схеме."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Zjhox1Rdmufa"
   },
   "source": [
    "### Политика"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "v5KOdiJrmufb"
   },
   "source": [
    "Наша политика должна принимать на вход состояние среды, а на выходе выдавать распределение по действиям, которые мы можем осуществлять в среде.\n",
    "\n",
    "**Задание:** Создать класс нейронной сети со следующей архитектурой ```Linear -> ReLU -> Linear -> Softmax```. Параметрами инициализации должны служить размерности пространства состояний, пространства действий и размер скрытого слоя."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BSbzinGQidFm"
   },
   "outputs": [],
   "source": [
    "class Policy(nn.Module):\n",
    "    \"\"\"Policy to be used by agent.\n",
    "\n",
    "    Attributes:\n",
    "        state_size: Dimention of the state space of the environment.\n",
    "        act_size: Dimention of the action space of the environment.\n",
    "        hidden_size: Dimention of the hidden state of the agent's policy.\n",
    "    \"\"\"\n",
    "    \n",
    "    # TO DO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "E0bYNhJ_mufh"
   },
   "source": [
    "### Оценка правдоподобия и расчет суммы дисконтированных наград"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QMZIoW8pmufi"
   },
   "source": [
    "**Задание:** Напишем вспомогательная функцию, которая принимает на вход политику, батч траекторий и фактор дисконтирования, и должна вернуть следующие величины: \n",
    "* правдоподобие действия на каждом шаге на траектории посчитанные для всего батча;\n",
    "* дисконтированные суммы наград (reward-to-go) из каждого состояния среды на траектории посчитанные для всего батча;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rcs_Xi-jmufl"
   },
   "source": [
    "**Hint**: Представим батч траекторий как ```list```, в котром также хранится ```list``` для каждой траектории, в котором каждый шаг хранится, как ```namedtuple```:\n",
    "```transition = collections.namedtuple(\"transition\", [\"state\", \"action\", \"reward\"])```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eGHfushNmufl"
   },
   "outputs": [],
   "source": [
    "def process_traj_batch(policy, batch, discount):\n",
    "    \"\"\"Computes log probabilities for each action \n",
    "        and rewards-to-go for each state in the batch of trajectories.\n",
    "    \n",
    "    Args:\n",
    "        policy: Policy of the agent.\n",
    "        batch (list of list of collections.namedtuple): Batch of trajectories.\n",
    "        discount (float): Discount factor for rewards-to-go calculation.\n",
    "\n",
    "    Returns:\n",
    "        log_probs (list of torch.FloatTensor): List of log probabilities for\n",
    "            each action in the batch of trajectories.\n",
    "        returns (list of rewards-to-go): List of rewards-to-go for \n",
    "            each state in the batch of trajectories.\n",
    "    \"\"\"\n",
    "    \n",
    "    # TO DO\n",
    "    \n",
    "    return log_probs, returns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mp6xKRWXmufp"
   },
   "source": [
    "Ваша реализация функции должна проходить следующий тест."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qVJnolh-mufq"
   },
   "outputs": [],
   "source": [
    "def test_process_traj_batch(process_traj_batch):\n",
    "    \n",
    "    transition = collections.namedtuple(\"transition\", [\"state\", \"action\", \"reward\"])\n",
    "    \n",
    "    class HelperPolicy(nn.Module):\n",
    "        def __init__(self):\n",
    "            super(HelperPolicy, self).__init__()\n",
    "        \n",
    "            self.act = nn.Sequential(\n",
    "                nn.Linear(4, 2),\n",
    "                nn.Softmax(dim=0),\n",
    "            )\n",
    "        \n",
    "        def forward(self, x):\n",
    "            return self.act(x)\n",
    "        \n",
    "    policy = HelperPolicy()\n",
    "    \n",
    "    for name, param in policy.named_parameters():\n",
    "        if name == \"act.0.weight\":\n",
    "            param.data = torch.tensor([[1.7492, -0.2471, 0.3310, 1.1494],\n",
    "                                       [0.6171, -0.6026, 0.5025, -0.3196]])\n",
    "        else:\n",
    "            param.data = torch.tensor([0.0262, 0.1882])\n",
    "            \n",
    "    batch = [\n",
    "        [\n",
    "            transition(state=torch.tensor([ 0.0462, -0.0018,  0.0372,  0.0063]), action=torch.tensor(0), reward=1.0), \n",
    "            transition(state=torch.tensor([ 0.0462, -0.1975,  0.0373,  0.3105]), action=torch.tensor(1), reward=1.0), \n",
    "            transition(state=torch.tensor([ 0.0422, -0.0029,  0.0435,  0.0298]), action=torch.tensor(0), reward=1.0), \n",
    "            transition(state=torch.tensor([ 0.0422, -0.1986,  0.0441,  0.3359]), action=torch.tensor(0), reward=1.0), \n",
    "        ],\n",
    "        [\n",
    "            transition(state=torch.tensor([ 0.0382, -0.3943,  0.0508,  0.6421]), action=torch.tensor(1), reward=1.0), \n",
    "            transition(state=torch.tensor([ 0.0303, -0.2000,  0.0637,  0.3659]), action=torch.tensor(1), reward=1.0), \n",
    "            transition(state=torch.tensor([ 0.0263, -0.0058,  0.0710,  0.0939]), action=torch.tensor(1), reward=1.0), \n",
    "            transition(state=torch.tensor([ 0.0262,  0.1882,  0.0729, -0.1755]), action=torch.tensor(0), reward=1.0)\n",
    "        ]\n",
    "    ]\n",
    "    \n",
    "    log_probs, returns = process_traj_batch(policy, batch, 0.9)\n",
    "    assert sum(log_probs).item() == -6.3940582275390625, \"Log probabilities calculation is incorrect!!!\"\n",
    "    assert sum(returns) == 18.098, \"Log probabilities calculation is incorrect!!!\"\n",
    "    print(\"Correct!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_qKrJpyBmufu"
   },
   "outputs": [],
   "source": [
    "test_process_traj_batch(process_traj_batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7_y4MJPxmufz"
   },
   "source": [
    "### Вспомогательные функции и гиперпараметры"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "PKF5itlpmuf0"
   },
   "source": [
    "Функция для расчета скользящего среднего - ее мы будем использовать для визуализации наград по эпизодам."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tHgXs9aPmuf2"
   },
   "outputs": [],
   "source": [
    "moving_average = lambda x, **kw: pd.DataFrame({'x':np.asarray(x)}).x.ewm(**kw).mean().values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "P3oqmQmxmuf6"
   },
   "source": [
    "Определим также гиперпараметры."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cFSLjGjOms0m"
   },
   "outputs": [],
   "source": [
    "STATE_SIZE  = env.observation_space.shape[0] # размерность пространства состояний среды\n",
    "ACT_SIZE = env.action_space.n # размерность пространства действий среды\n",
    "HIDDEN_SIZE = 256 # размер скрытого слоя для политики\n",
    "NUM_EPISODES = 1000 # количество эпиздов, которые будут сыграны для обучения\n",
    "DISCOUNT = 0.99 # фактор дисконтирования\n",
    "TRAIN_EVERY = 20"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "d1RZfo_1muf-"
   },
   "source": [
    "Инициализуем политику и алгоритм оптимизации - мы будем использовать Adam c праметрами по умолчанию."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zEKqg4Ykmuf_"
   },
   "outputs": [],
   "source": [
    "policy = Policy(STATE_SIZE, ACT_SIZE, HIDDEN_SIZE).to(device)\n",
    "optimizer = Adam(policy.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EIbRMh3MmugC"
   },
   "outputs": [],
   "source": [
    "transition = collections.namedtuple(\"transition\", [\"state\", \"action\", \"reward\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jRRCJTnimugE"
   },
   "source": [
    "### Основной цикл обучения"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FcQ4-JMymugE"
   },
   "source": [
    "Теперь, когда мы опредлели вспомогательные функции, то нам следует написать основной цикл обучения агент.\n",
    "В цикле должно происходить следующее:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1wih2vIYmugG"
   },
   "source": [
    "1. Играем количество эпизодов, определенное в гиперпараметре ```NUM_EPISODES```.\n",
    "2. В каждом эпизоде сохраняем информацию о шагах на траектории - состояние, действие и награду.\n",
    "3. В конце каждого эпизода сохраняем вышеуказанную информацию о траектории. \n",
    "4. Периодически обучаемся на собранных эпизодах каждые ```TRAIN_EVERY``` эпизодов:  \n",
    "    4.1. Считаем для собранного батча для каждого шага на трактории правдоподобие и сумму дисконтированных наград.  \n",
    "    4.2. Обновляем параметры политики агента по формуле, приведенной на схеме."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gb1eIuJqmugI"
   },
   "source": [
    "**Задание:** Реализовать алгоритм обучения, описанный на схеме и в тексте выше. Шаблон кода алгоритма представлен ниже. При этом следует сохранять сумму ревордов для каждого эпизода в переменную ```returns_history```. Алгоритму потребуется около 1000 эпизодов игры, для того чтобы научиться играть в игру (если после 1000 эпизодов агент играет немного хуже, чем для победы в игре, попробуйте обучать его немного дольше или установите критерий останова - когда средняя награда за 100 последних эпизодов превышает значение в ```env.spec.reward_threshold``` )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cvEwjezan5tl"
   },
   "outputs": [],
   "source": [
    "returns_history = []\n",
    "traj_batch = []\n",
    "\n",
    "for i in range(NUM_EPISODES):\n",
    "        \n",
    "    # TO DO\n",
    "        \n",
    "    returns_history.append(rewards)\n",
    "            \n",
    "    traj_batch.append(traj)\n",
    "    \n",
    "    if i % TRAIN_EVERY:\n",
    "        log_probs, returns = process_traj_batch(policy, traj_batch, DISCOUNT)\n",
    "        loss = -(torch.stack(log_probs) * torch.FloatTensor(returns).to(device)).sum()\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        traj_batch = []\n",
    "\n",
    "    if i % 10:\n",
    "        clear_output(True)\n",
    "        plt.figure(figsize=[12, 6])\n",
    "        plt.title('Returns'); plt.grid()\n",
    "        plt.scatter(np.arange(len(returns_history)), returns_history, alpha=0.1)\n",
    "        plt.plot(moving_average(returns_history, span=10, min_periods=10))\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Be9Zh-MamugN"
   },
   "source": [
    "Протестируем обученного агента."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-ry9g_aBmugO"
   },
   "outputs": [],
   "source": [
    "test_agent(env, agent=policy, n_episodes=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rf18pOwCmugS"
   },
   "source": [
    "Обученный агент должен приближаться к искомому значению средней награды за 100 эпизодов 195."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Gk3bQdF_mugT"
   },
   "source": [
    "Визуализируем обученного агента."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "c6QvXvwUsct6"
   },
   "outputs": [],
   "source": [
    "agent_viz(agent=policy)\n",
    "show_video()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "YKkrymnnmugU"
   },
   "source": [
    "Как видно, агент выучил довольно хорошую стратегию для игры и способен долго удерживать шест."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NHxdzJsQmugV"
   },
   "source": [
    "### REINFORCE with baselines (Опционально)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QM0DmpwlmugX"
   },
   "source": [
    "В лекциях вы слышали, что при расчете градиентов для обновления параметров политики агента мы можем вычесть из суммы дисконтированных наград ```baseline``` для уменьшения дисперсии градиентов и ускорения сходимости обучения - такой алгоритм называется REINFORCE with baselines. В качестве ```baseline``` мы можем использовать другую нейронную сеть, которая будет оценивать сумму дисконтированных наград из данного состояния *V(s)*. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jImFlpefmugX"
   },
   "source": [
    "Схема алгоритма REINFORCE with baselines (Sutton & Barto) <img src=\"//i.imgur.com/j3BcbHP.png\" width=\"700\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ofDdKiKJmugY"
   },
   "source": [
    "**Задание**: Включите в уже разработанный алгоритм вторую нейронную сеть для оценки суммы дисконтированных наград *V(s)*. Используйте разницу между фактической суммой дисконтированных наград и оценкой в формуле функции потерь политики. В качестве функции потерь для *V(s)* используйте ```MSELoss```. Оцените скорость сходимости нового алгоритма."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "REINFORCE.ipynb",
   "private_outputs": true,
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
